package MLDA_Process;

import gnu.trove.map.hash.TIntIntHashMap;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import org.netlib.util.doubleW;

import Commons.IOUtil;

public class SparseSATM implements Runnable {

    public int P = 200;
    public int K = 20;

    public double alpha = 50.0 / K;
    public double beta = 0.1;

    public int niters = 100;
    public int save_step = 500;

    public int M;
    public int V;

    public double kAlpha;
    public double vBeta;

    TIntIntHashMap[] npkTMaps;
    TIntIntHashMap[] nwkTMaps;

    public int[] nkwSum;
    public int[] npkSum;

    public int npw[][]; // P * V
    public int npwSum[]; // P * 1

    public double corpus_size;

    public double[][] eta;

    public int[][] z_assigns;
    public int[][] d_assigns;

    public double epsilon = 0.001;

    public double[][] c;

    public double e_sum;
    public double f_sum;
    public double g_sum;

    public String inputPath = "";
    public String outputPath = "";

    public List<List<Integer>> docs = new ArrayList<List<Integer>>();
    public HashMap<String, Integer> w2i = new HashMap<String, Integer>();
    public HashMap<Integer, String> i2w = new HashMap<Integer, String>();

    public SparseSATM(int p, int k, double alpha, double beta, int niters, int saveStep, String inputPath, String outputPath) {
        this.P = p;
        this.K = k;
        this.alpha = alpha;
        this.beta = beta;
        this.niters = niters;
        this.save_step = saveStep;
        this.inputPath = inputPath;
        this.outputPath = outputPath;
    }

    public void load_txts(String txtPath) {
        BufferedReader reader = IOUtil.getReader(txtPath, "UTF-8");

        String line;
        try {
            line = reader.readLine();
            while (line != null) {
                List<Integer> doc = new ArrayList<Integer>();

                String[] tokens = line.trim().split("\\s+");
                for (String token : tokens) {
                    if (!w2i.containsKey(token)) {
                        w2i.put(token, w2i.size());
                        i2w.put(w2i.get(token), token);
                    }
                    doc.add(w2i.get(token));
                }
                docs.add(doc);
                line = reader.readLine();
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        M = docs.size();
        V = w2i.size();
        return;
    }

    public void init_model() {

        corpus_size = 0.0;

        c = new double[P][K];

        vBeta = V * beta;
        kAlpha = K * alpha;

        nkwSum = new int[K];
        npkSum = new int[P];

        npw = new int[V][P];
        npwSum = new int[P];

        nwkTMaps = new TIntIntHashMap[V];
        for (int w = 0; w != V; w++) {
            nwkTMaps[w] = new TIntIntHashMap();
        }

        npkTMaps = new TIntIntHashMap[P];
        for (int d = 0; d != P; d++) {
            npkTMaps[d] = new TIntIntHashMap();
        }

        z_assigns = new int[M][];
        d_assigns = new int[M][];

        for (int m = 0; m != M; m++) {
            int N = docs.get(m).size();
            corpus_size += N;

            z_assigns[m] = new int[N];
            d_assigns[m] = new int[N];

            int d = (int) Math.floor(Math.random() * P);
            TIntIntHashMap npkTMap = npkTMaps[d];
            for (int n = 0; n != N; n++) {
                int w = docs.get(m).get(n);
                int z = (int) Math.floor(Math.random() * K);

                TIntIntHashMap nwkTMap = nwkTMaps[w];

                npw[w][d]++;
                npwSum[d]++;

                npkTMap.adjustOrPutValue(z, 1, 1);
                nwkTMap.adjustOrPutValue(z, 1, 1);

                nwkTMaps[w] = nwkTMap;

                npkSum[d]++;
                nkwSum[z]++;

                z_assigns[m][n] = z;
                d_assigns[m][n] = d;

            }
            npkTMaps[d] = npkTMap;
        }

        update_eta();
    }

    public void update_eta() {
        eta = new double[M][P];

        for (int m = 0; m != M; m++) {
            double norm = 0.0;
            for (int p = 0; p != P; p++) {
                eta[m][p] = (npwSum[p] + 1.0) / (corpus_size + P);
                for (int n = 0; n != docs.get(m).size(); n++) {
                    int w = docs.get(m).get(n);
                    eta[m][p] *= (npw[w][p] + 1.0) / ((double) npwSum[p] + V);
                }
                norm += eta[m][p];
            }
            for (int p = 0; p != P; p++) {
                eta[m][p] /= norm;
            }
        }
    }

    public void sample_doc(int m) {
        for (int j = 0; j != P; j++) {
            if (eta[m][j] < epsilon) continue;
            for (int k = 0; k != K; k++) {
                c[j][k] = alpha / ((nkwSum[k] + vBeta) * (npkSum[j] + kAlpha));
            }
        }

        g_sum = 0.0;
        for (int j = 0; j != P; j++) {
            if (eta[m][j] < epsilon) continue;
            for (int k = 0; k != K; k++) {
                g_sum += eta[m][j] * beta * c[j][k];
            }
        }

        f_sum = 0.0;
        for (int j = 0; j != P; j++) {
            if (eta[m][j] < epsilon) continue;
            TIntIntHashMap npkTMap = npkTMaps[j];
            for (int k : npkTMap.keys()) {
                f_sum += eta[m][j] * beta * npkTMap.get(k) / ((nkwSum[k] + vBeta) * (npkSum[j] + kAlpha));
            }
            for (int k = 0; k != K; k++) {
                c[j][k] += npkTMap.get(k) / ((nkwSum[k] + vBeta) * (npkSum[j] + kAlpha));
            }
        }

        for (int n = 0; n != docs.get(m).size(); n++) {
            int d = d_assigns[m][n];
            int z = z_assigns[m][n];

            int old_d = d;
            int old_z = z;
            int w = docs.get(m).get(n);

            TIntIntHashMap npkTMap = npkTMaps[old_d];
            TIntIntHashMap nwkTMap = nwkTMaps[w];

            double denom = (npkSum[old_d] + kAlpha) * (nkwSum[old_z] + vBeta);
            f_sum -= eta[m][old_d] * beta * npkTMap.get(old_z) / denom;
            g_sum -= (alpha * eta[m][old_d] * beta) / denom;

            if (npkTMap.get(old_z) == 1) {
                npkTMap.remove(old_z);
            } else {
                npkTMap.adjustValue(old_z, -1);
            }
            npkTMaps[old_d] = npkTMap;

            if (nwkTMap.get(old_z) == 1) {
                nwkTMap.remove(old_z);
            } else {
                nwkTMap.adjustValue(old_z, -1);
            }

            npkSum[old_d]--;
            nkwSum[old_z]--;

            npw[w][old_d]--;
            npwSum[old_d]--;

            denom = (npkSum[old_d] + kAlpha) * (nkwSum[old_z] + vBeta);

            if (npkTMap.containsKey(old_z)) {
                c[old_d][old_z] = (alpha + npkTMap.get(old_z)) / denom;
                f_sum += eta[m][old_d] * beta * npkTMap.get(old_z) / denom;
            } else {
                c[old_d][old_z] = alpha / denom;
//				f_sum += 0.0;
            }

            g_sum += (alpha * eta[m][d] * beta) / denom;

            e_sum = 0.0;
            for (int j = 0; j != P; j++) {
                for (int k : nwkTMap.keys()) {
                    e_sum += nwkTMap.get(k) * eta[m][j] * c[j][k];
                }
            }

            boolean end = false;
            double total_mass = e_sum + f_sum + g_sum;
            double u = Math.random() * total_mass;
            double r = 0.0, pSum = 0.0;
            if (u <= e_sum) {
                r = Math.random() * e_sum;
                for (int j = 0; j != P; j++) {
                    if (eta[m][j] < epsilon) continue;
                    for (int k : nwkTMap.keys()) {
                        pSum += eta[m][j] * nwkTMap.get(k) * c[j][k];
                        if (r < pSum) {
                            d = j;
                            z = k;
                            end = true;
                            break;
                        }
                    }
                    if (end) break;
                }
            } else {
                if (u <= e_sum + f_sum) {
                    r = Math.random() * f_sum;
                    for (int j = 0; j != P; j++) {
                        if (eta[m][j] < epsilon) continue;
                        npkTMap = npkTMaps[j];
                        for (int k : npkTMap.keys()) {
                            pSum += eta[m][j] * beta * npkTMap.get(k) / ((npkSum[j] + kAlpha) * (nkwSum[k] + vBeta));
                            if (r < pSum) {
                                d = j;
                                z = k;
                                end = true;
                                break;
                            }
                        }
                        if (end) break;
                    }
                } else {
                    r = Math.random() * g_sum;
                    for (int j = 0; j != P; j++) {
                        if (eta[m][j] < epsilon) continue;
                        for (int k = 0; k != K; k++) {
                            pSum += eta[m][j] * beta * alpha / ((npkSum[j] + kAlpha) * (nkwSum[k] + vBeta));
                            if (r < pSum) {
                                z = k;
                                d = j;
                                end = true;
                                break;
                            }
                        }
                        if (end) break;
                    }
                }
            }

            npkTMap = npkTMaps[d];

            denom = (npkSum[d] + kAlpha) * (nkwSum[z] + vBeta);
            if (npkTMap.containsKey(z))
                f_sum -= eta[m][d] * beta * npkTMap.get(z) / denom;
            g_sum -= (alpha * eta[m][d] * beta) / denom;

            nwkTMap.adjustOrPutValue(z, 1, 1);
            npkTMap.adjustOrPutValue(z, 1, 1);

            nwkTMaps[w] = nwkTMap;
            npkTMaps[d] = npkTMap;

            nkwSum[z]++;
            npkSum[d]++;

            npw[w][d]++;
            npwSum[d]++;

            denom = (npkSum[d] + kAlpha) * (nkwSum[z] + vBeta);

            f_sum += (eta[m][d] * beta * npkTMap.get(z)) / denom;

            g_sum += (alpha * eta[m][d] * beta) / denom;

            c[d][z] = (alpha + npkTMap.get(z)) / denom;

            z_assigns[m][n] = z;
            d_assigns[m][n] = d;
        }
    }

    int nkw[][];

    public void convert_map_to_arrays() {
        nkw = new int[K][V];
        for (int w = 0; w != V; w++) {
            TIntIntHashMap nwkTMap = nwkTMaps[w];
            for (int k : nwkTMap.keys()) {
                nkw[k][w] = nwkTMap.get(k);
            }
        }
    }

    public double[][] computePhi() {
        convert_map_to_arrays();
        double[][] phi = new double[K][V];
        for (int k = 0; k != K; k++) {
            for (int v = 0; v != V; v++) {
                phi[k][v] = (nkw[k][v] + beta) / (nkwSum[k] + V * beta);
            }
        }
        return phi;
    }

    int ndk[][];
    int ndkSum[];

    public void convert_zassigns_to_arrays_theta() {
        ndk = new int[M][K];
        ndkSum = new int[M];

        for (int m = 0; m != M; m++) {
            for (int n = 0; n != docs.get(m).size(); n++) {
                ndk[m][z_assigns[m][n]]++;
                ndkSum[m]++;
            }
        }
    }

    public double[][] computeTheta() {
        convert_zassigns_to_arrays_theta();
        double[][] theta = new double[M][K];
        for (int m = 0; m != M; m++) {
            for (int k = 0; k != K; k++) {
                theta[m][k] = (ndk[m][k] + alpha) / (ndkSum[m] + K * alpha);
            }
        }
        return theta;
    }

    public void saveTheta(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] theta = computeTheta();
        try {
            for (int m = 0; m != M; m++) {
                for (int k = 0; k != K; k++) {
                    writer.append(theta[m][k] + " ");
                }
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public ArrayList<List<Entry<String, Double>>> sortedTopicWords(
            double[][] phi, int T) {
        ArrayList<List<Entry<String, Double>>> res = new ArrayList<List<Entry<String, Double>>>();
        for (int k = 0; k != T; k++) {
            HashMap<String, Double> term2weight = new HashMap<String, Double>();
            for (String term : w2i.keySet())
                term2weight.put(term, phi[k][w2i.get(term)]);

            List<Entry<String, Double>> pairs = new ArrayList<Entry<String, Double>>(
                    term2weight.entrySet());
            Collections.sort(pairs, new Comparator<Entry<String, Double>>() {
                public int compare(Entry<String, Double> o1,
                                   Entry<String, Double> o2) {
                    return (o2.getValue().compareTo(o1.getValue()));
                }
            });
            res.add(pairs);
        }
        return res;
    }

    public void printTopics(String path, int top_n) throws IOException {
        double[][] phi = computePhi();
        BufferedWriter writer = IOUtil.getWriter(path);
        ArrayList<List<Entry<String, Double>>> pairsList = this
                .sortedTopicWords(phi, K);
        for (int k = 0; k != K; k++) {
            writer.write("Topic " + k + ":\n");
            for (int i = 0; i != top_n; i++) {
                writer.write(pairsList.get(k).get(i).getKey() + " "
                        + pairsList.get(k).get(i).getValue() + "\n");
            }
        }
        writer.close();
    }

    public void savePhi(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] phi = computePhi();
        int K = phi.length;
        assert K > 0;
        int V = phi[0].length;

        try {
            for (int k = 0; k != K; k++) {
                for (int v = 0; v != V; v++) {
                    writer.append(phi[k][v] + " ");
                }
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveWordmap(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        try {
            for (String word : w2i.keySet())
                writer.append(word + "\t" + w2i.get(word) + "\n");

            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void estimate() {
        for (int iter = 0; iter != niters; iter++) {
            System.out.println("Sparse SATM Iteration " + iter + " ...");
            if (iter % this.save_step == 0 && iter != 0) {
                this.storeResult(iter);
            }

            long begin = System.currentTimeMillis();
            for (int m = 0; m != M; m++) {
                sample_doc(m);
            }
            System.out.println((System.currentTimeMillis() - begin) / 1000);
            this.update_eta();
            System.out.println((System.currentTimeMillis() - begin) / 1000);
        }
    }

    public void saveAssign(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");
        try {
            for (int i = 0; i < z_assigns.length; i++) {
                for (int j = 0; j < z_assigns[i].length; j++) {
                    writer.write(docs.get(i).get(j) + ":" + z_assigns[i][j] + " ");
                }
                writer.write("\n");
            }

            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        return;
    }

    public void printModel() {
        System.out.println("\tP :" + this.P +
                "\tK :" + this.K +
                "\tniters :" + this.niters +
                "\tsaveStep :" + this.save_step +
                "\talpha :" + this.alpha +
                "\tbeta :" + this.beta +
                "\tinputPath :" + this.inputPath +
                "\toutputPath :" + this.outputPath);
    }

    public void storeResult(int times) {
        String appendString = "final";
        if (times != 0) {
            appendString = times + "";
        }
        try {
            this.printTopics(outputPath + "/model-" + appendString + ".twords", 20);
            this.saveWordmap(outputPath + "/wordmap.txt");
            this.savePhi(outputPath + "/model-" + appendString + ".phi");
            this.saveAssign(outputPath + "/model-" + appendString + ".tassign");
            this.saveTheta(outputPath + "/model-" + appendString + ".theta");
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }

    public void run() {
        printModel();
        this.load_txts(inputPath);//
        this.init_model();
        this.estimate();
        this.convert_map_to_arrays();
        storeResult(0);

    }

    public static void SATMForShortText(int k1, int k2, double alpha, double beta, int iter, int saveStep, int threadNum, String path) {
        ExecutorService service = Executors.newCachedThreadPool();
        File trainFile = new File(path);
        String parent_path = trainFile.getParentFile().getAbsolutePath()
                + "/SP_SATM_" + k1 + "_" + k2 + "_" + alpha + "_" + beta + "/";
        for (int i = 0; i < threadNum; i++) {
            (new File(parent_path + i)).mkdirs();
            service.submit(new SparseSATM(k1, k2, alpha, beta, iter, saveStep, path, parent_path + i + "/"));

        }
        service.shutdown();
    }

}
